{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 解读 `tvm.ir.structural_hash`\n",
    "\n",
    "参考：`ai/tvm/tests/python/relay/test_ir_structural_equal_hash.py`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "from pathlib import Path\n",
    "ROOT = Path(\".\").resolve().parents[2]\n",
    "sys.path.extend([f\"{ROOT}/tests\", f\"{ROOT}/src\"])\n",
    "# # from tools.tag_span import _create_span, _set_span, _verify_structural_equal_with_span\n",
    "from tools.torch_utils import verify_model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import tvm\n",
    "import tvm.testing\n",
    "from tvm import relay\n",
    "from tvm.relay.testing import run_opt_pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def consistent_equal(x, y, map_free_vars=False):\n",
    "    struct_equal0 = tvm.ir.structural_equal(x, y, map_free_vars)\n",
    "    struct_equal1 = tvm.ir.structural_equal(y, x, map_free_vars)\n",
    "\n",
    "    xhash = tvm.ir.structural_hash(x, map_free_vars)\n",
    "    yhash = tvm.ir.structural_hash(y, map_free_vars)\n",
    "\n",
    "    if struct_equal0 != struct_equal1:\n",
    "        raise ValueError(\n",
    "            \"Non-communicative {} vs {}, sequal0={}, sequal1={}\".format(\n",
    "                x, y, struct_equal0, struct_equal1\n",
    "            )\n",
    "        )\n",
    "\n",
    "    # NOTE: hash colision can happen but should be rare.\n",
    "    # we can confirm that hash colison doesn't happen for our testcases\n",
    "    if struct_equal0 != (xhash == yhash):\n",
    "        raise ValueError(\n",
    "            \"Inconsistent {} vs {}, sequal={}, xhash={}, yhash={}\".format(\n",
    "                x, y, struct_equal0, xhash, yhash\n",
    "            )\n",
    "        )\n",
    "    return struct_equal0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_tensor_type_sequal():\n",
    "    t1 = relay.TensorType((3, 4), \"float32\")\n",
    "    t2 = relay.TensorType((3, 4), \"float32\")\n",
    "    t3 = relay.TensorType((3, 4, 5), \"float32\")\n",
    "    assert t1 == t2\n",
    "    assert t1 != t3\n",
    "\n",
    "    t1 = relay.TensorType((), \"float32\")\n",
    "    t2 = relay.TensorType((), \"float32\")\n",
    "    assert t1 == t2\n",
    "\n",
    "\n",
    "def test_incomplete_type_sequal():\n",
    "    t1 = relay.IncompleteType(relay.TypeKind.ShapeVar)\n",
    "    t2 = relay.IncompleteType(relay.TypeKind.Type)\n",
    "    t3 = relay.IncompleteType(relay.TypeKind.Type)\n",
    "\n",
    "    # only equal when there is pointer equality\n",
    "    assert t2 == t2\n",
    "    assert t1 == t1\n",
    "    assert t1 != t2\n",
    "    assert t2 != t3\n",
    "\n",
    "\n",
    "def test_type_param_sequal():\n",
    "    t1 = relay.TypeVar(\"v1\", relay.TypeKind.Type)\n",
    "    t2 = relay.TypeVar(\"v2\", relay.TypeKind.ShapeVar)\n",
    "    t3 = relay.TypeVar(\"v3\", relay.TypeKind.Type)\n",
    "\n",
    "    # only pointer equality and eq_map allow equal params\n",
    "    assert t1 == t1\n",
    "    assert t2 == t2\n",
    "    assert t1 != t2  # different kind\n",
    "    assert t1 != t3  # not in eq_map\n",
    "\n",
    "    # function types are the only way to put type params\n",
    "    # in eq map\n",
    "    ft1 = relay.FuncType(\n",
    "        tvm.runtime.convert([]), t1, tvm.runtime.convert([t1]), tvm.runtime.convert([])\n",
    "    )\n",
    "    ft2 = relay.FuncType(\n",
    "        tvm.runtime.convert([]), t3, tvm.runtime.convert([t3]), tvm.runtime.convert([])\n",
    "    )\n",
    "    # actually an invalid type because t2 is wrong kind\n",
    "    ft3 = relay.FuncType(\n",
    "        tvm.runtime.convert([]), t2, tvm.runtime.convert([t2]), tvm.runtime.convert([])\n",
    "    )\n",
    "\n",
    "    assert ft1 == ft2\n",
    "    assert ft1 != ft3  # kinds still do not match\n",
    "\n",
    "\n",
    "def test_func_type_sequal():\n",
    "    t1 = relay.TensorType((1, 2), \"float32\")\n",
    "    t2 = relay.TensorType((1, 2, 3), \"float32\")\n",
    "\n",
    "    tp1 = relay.TypeVar(\"v1\", relay.TypeKind.Type)\n",
    "    tp2 = relay.TypeVar(\"v2\", relay.TypeKind.Type)\n",
    "    tp3 = relay.TypeVar(\"v3\", relay.TypeKind.ShapeVar)\n",
    "    tp4 = relay.TypeVar(\"v3\", relay.TypeKind.ShapeVar)\n",
    "\n",
    "    broadcast = tvm.ir.EnvFunc.get(\"tvm.relay.type_relation.Broadcast\")\n",
    "    identity = tvm.ir.EnvFunc.get(\"tvm.relay.type_relation.Identity\")\n",
    "\n",
    "    tr1 = relay.TypeRelation(broadcast, tvm.runtime.convert([tp1, tp3]), 1, None)\n",
    "    tr2 = relay.TypeRelation(broadcast, tvm.runtime.convert([tp2, tp4]), 1, None)\n",
    "    tr3 = relay.TypeRelation(identity, tvm.runtime.convert([tp1, tp3]), 1, None)\n",
    "\n",
    "    ft = relay.FuncType(\n",
    "        tvm.runtime.convert([t1, t2]),\n",
    "        tp1,\n",
    "        tvm.runtime.convert([tp1, tp3]),\n",
    "        tvm.runtime.convert([tr1]),\n",
    "    )\n",
    "    translate_vars = relay.FuncType(\n",
    "        tvm.runtime.convert([t1, t2]),\n",
    "        tp2,\n",
    "        tvm.runtime.convert([tp2, tp4]),\n",
    "        tvm.runtime.convert([tr2]),\n",
    "    )\n",
    "    assert ft == translate_vars\n",
    "\n",
    "    different_args = relay.FuncType(\n",
    "        tvm.runtime.convert([t1]), tp1, tvm.runtime.convert([tp1, tp3]), tvm.runtime.convert([tr1])\n",
    "    )\n",
    "    assert ft != different_args\n",
    "\n",
    "    different_order = relay.FuncType(\n",
    "        tvm.runtime.convert([t2, t1]),\n",
    "        tp1,\n",
    "        tvm.runtime.convert([tp1, tp3]),\n",
    "        tvm.runtime.convert([tr1]),\n",
    "    )\n",
    "    assert ft != different_order\n",
    "\n",
    "    no_rel = relay.FuncType(\n",
    "        tvm.runtime.convert([t1, t2]), tp1, tvm.runtime.convert([tp1, tp3]), tvm.runtime.convert([])\n",
    "    )\n",
    "    assert ft != no_rel\n",
    "\n",
    "    more_vars = relay.FuncType(\n",
    "        tvm.runtime.convert([t1, t2]),\n",
    "        tp2,\n",
    "        tvm.runtime.convert([tp1, tp2, tp3]),\n",
    "        tvm.runtime.convert([tr1]),\n",
    "    )\n",
    "    assert ft != more_vars\n",
    "\n",
    "    all_the_vars = relay.FuncType(\n",
    "        tvm.runtime.convert([t1, t2]),\n",
    "        tp1,\n",
    "        tvm.runtime.convert([tp1, tp2, tp3, tp4]),\n",
    "        tvm.runtime.convert([tr1, tr2]),\n",
    "    )\n",
    "    assert ft != all_the_vars\n",
    "\n",
    "    different_rel = relay.FuncType(\n",
    "        tvm.runtime.convert([t1, t2]),\n",
    "        tp1,\n",
    "        tvm.runtime.convert([tp1, tp3]),\n",
    "        tvm.runtime.convert([tr3]),\n",
    "    )\n",
    "    assert ft != different_rel\n",
    "\n",
    "    more_rels = relay.FuncType(\n",
    "        tvm.runtime.convert([t1, t2]),\n",
    "        tp1,\n",
    "        tvm.runtime.convert([tp1, tp3]),\n",
    "        tvm.runtime.convert([tr1, tr3]),\n",
    "    )\n",
    "    assert ft != more_rels\n",
    "\n",
    "\n",
    "def test_tuple_type_sequal():\n",
    "    t1 = relay.TensorType((1, 2, 3), \"float32\")\n",
    "    t2 = relay.TensorType((1, 2, 3, 4), \"float32\")\n",
    "    tp1 = relay.TypeVar(\"v1\", relay.TypeKind.Type)\n",
    "    tp2 = relay.TypeVar(\"v2\", relay.TypeKind.Type)\n",
    "\n",
    "    tup1 = relay.TupleType(tvm.runtime.convert([t1, t2, tp1]))\n",
    "    tup2 = relay.TupleType(tvm.runtime.convert([t1, t2, tp1]))\n",
    "    tup3 = relay.TupleType(tvm.runtime.convert([t2, t1, tp1]))\n",
    "    tup4 = relay.TupleType(tvm.runtime.convert([t1, t2, tp2]))\n",
    "\n",
    "    # as long as types are alpha-equal and in same order,\n",
    "    # tuples should be alpha-equal\n",
    "    assert tup1 == tup2\n",
    "    assert tup1 != tup3\n",
    "    assert tup1 != tup4\n",
    "\n",
    "\n",
    "def test_type_relation_sequal():\n",
    "    t1 = relay.TensorType((1, 2), \"float32\")\n",
    "    t2 = relay.TensorType((1, 2, 3), \"float32\")\n",
    "    t3 = relay.TensorType((1, 2, 3, 4), \"float32\")\n",
    "\n",
    "    # functions are compared only by pointer equality so\n",
    "    # we need to be sure to use the same pointers\n",
    "    broadcast = tvm.ir.EnvFunc.get(\"tvm.relay.type_relation.Broadcast\")\n",
    "    identity = tvm.ir.EnvFunc.get(\"tvm.relay.type_relation.Identity\")\n",
    "\n",
    "    attr1 = tvm.ir.make_node(\"attrs.TestAttrs\", name=\"attr\", padding=(3, 4))\n",
    "    attr1_same = tvm.ir.make_node(\"attrs.TestAttrs\", name=\"attr\", padding=(3, 4))\n",
    "    attr2 = tvm.ir.make_node(\"attrs.TestAttrs\", name=\"attr\", padding=(3, 4, 4))\n",
    "\n",
    "    tr = relay.TypeRelation(broadcast, tvm.runtime.convert([t1, t2]), 1, attr1)\n",
    "    same = relay.TypeRelation(broadcast, tvm.runtime.convert([t1, t2]), 1, attr1)\n",
    "    diff_func = relay.TypeRelation(identity, tvm.runtime.convert([t1, t2]), 1, attr1)\n",
    "    diff_order = relay.TypeRelation(broadcast, tvm.runtime.convert([t2, t1]), 1, attr1)\n",
    "    diff_args = relay.TypeRelation(broadcast, tvm.runtime.convert([t2, t3]), 1, attr1)\n",
    "    diff_attr = relay.TypeRelation(broadcast, tvm.runtime.convert([t1, t2]), 1, attr2)\n",
    "    same_attr = relay.TypeRelation(broadcast, tvm.runtime.convert([t1, t2]), 1, attr1_same)\n",
    "\n",
    "    bigger = relay.TypeRelation(identity, tvm.runtime.convert([t1, t3, t2]), 2, attr1)\n",
    "    diff_num_inputs = relay.TypeRelation(identity, tvm.runtime.convert([t1, t3, t2]), 1, attr2)\n",
    "\n",
    "    # func, number of args, input count, and order should be the same\n",
    "    assert tr == same\n",
    "    assert tr != diff_func\n",
    "    assert tr != diff_order\n",
    "    assert tr != diff_args\n",
    "    assert tr != diff_attr\n",
    "    assert tr == same_attr\n",
    "    assert tr != bigger\n",
    "\n",
    "    assert bigger != diff_num_inputs\n",
    "\n",
    "\n",
    "def test_type_call_sequal():\n",
    "    h1 = relay.GlobalTypeVar(\"h1\")\n",
    "    h2 = relay.GlobalTypeVar(\"h2\")\n",
    "    t1 = relay.TensorType((1, 2), \"float32\")\n",
    "    t2 = relay.TensorType((1, 2, 3), \"float32\")\n",
    "    t3 = relay.TensorType((1, 2, 3, 4), \"float32\")\n",
    "    t4 = relay.TensorType((), \"float32\")\n",
    "\n",
    "    tc = relay.TypeCall(h1, [t1, t2, t3])\n",
    "    same = relay.TypeCall(h1, [t1, t2, t3])\n",
    "\n",
    "    different_func = relay.TypeCall(h2, [t1, t2, t3])\n",
    "    different_arg = relay.TypeCall(h1, [t1, t2, t4])\n",
    "    fewer_args = relay.TypeCall(h1, [t1, t2])\n",
    "    more_args = relay.TypeCall(h1, [t1, t2, t3, t4])\n",
    "    different_order_args = relay.TypeCall(h1, [t3, t2, t1])\n",
    "\n",
    "    assert tc == same\n",
    "    assert tc != different_func\n",
    "    assert tc != fewer_args\n",
    "    assert tc != more_args\n",
    "    assert tc != different_order_args\n",
    "\n",
    "\n",
    "def test_constant_sequal():\n",
    "    x = relay.const(1)\n",
    "    y = relay.const(2)\n",
    "    assert consistent_equal(x, x)\n",
    "    assert not consistent_equal(x, y)\n",
    "    assert consistent_equal(x, relay.const(1))\n",
    "\n",
    "\n",
    "def test_type_node_sequal():\n",
    "    v1 = relay.TypeVar(\"v1\", 6)\n",
    "    v2 = relay.TypeVar(\"v2\", 6)\n",
    "    assert not consistent_equal(v1, v2)\n",
    "\n",
    "    v1 = relay.TypeVar(\"v1\", 0)\n",
    "    v2 = relay.TypeVar(\"v2\", 6)\n",
    "    assert not consistent_equal(v1, v2)\n",
    "\n",
    "\n",
    "def test_type_node_incompatible_sequal():\n",
    "    v1 = relay.TypeVar(\"v1\", 6)\n",
    "    v2 = relay.Var(\"v2\")\n",
    "    assert not consistent_equal(v1, v2)\n",
    "\n",
    "\n",
    "def test_expr_node_incompatible_sequal():\n",
    "    v1 = relay.Var(\"v1\")\n",
    "    v2 = relay.PatternVar(relay.Var(\"v2\"))\n",
    "    assert not consistent_equal(v1, v2)\n",
    "\n",
    "\n",
    "def test_var_sequal():\n",
    "    v1 = relay.Var(\"v1\")\n",
    "    v2 = relay.Var(\"v2\")\n",
    "\n",
    "    # normally only pointer equality\n",
    "    assert consistent_equal(v1, v1)\n",
    "    assert not consistent_equal(v1, v2)\n",
    "\n",
    "    # let node allows for setting the eq_map\n",
    "    l1 = relay.Let(v1, relay.const(1), v1)\n",
    "    l2 = relay.Let(v2, relay.const(1), v2)\n",
    "    l3 = relay.Let(v1, relay.const(1), v2)\n",
    "\n",
    "    assert consistent_equal(l1, l2)\n",
    "    assert not consistent_equal(l1, l3)\n",
    "\n",
    "    # type annotations\n",
    "    tt1 = relay.TensorType([], \"int32\")\n",
    "    tt2 = relay.TensorType([], \"int32\")\n",
    "    tt3 = relay.TensorType([], \"int64\")\n",
    "    v3 = relay.Var(\"v3\", tt1)\n",
    "    v4 = relay.Var(\"v4\", tt2)\n",
    "    v5 = relay.Var(\"v5\", tt3)\n",
    "\n",
    "    l4 = relay.Let(v3, relay.const(1), v3)\n",
    "    l5 = relay.Let(v4, relay.const(1), v4)\n",
    "    l6 = relay.Let(v5, relay.const(1), v5)\n",
    "\n",
    "    # same annotations\n",
    "    assert consistent_equal(l4, l5)\n",
    "    # different annotations\n",
    "    assert not consistent_equal(l4, l6)\n",
    "    # one null annotation\n",
    "    assert not consistent_equal(l1, l4)\n",
    "\n",
    "\n",
    "def test_global_var_sequal():\n",
    "    v1 = relay.GlobalVar(\"v1\")\n",
    "    v2 = relay.GlobalVar(\"v2\")\n",
    "\n",
    "    # only pointer equality suffices (smoke test)\n",
    "    assert consistent_equal(v1, v1)\n",
    "    assert not consistent_equal(v1, v2)\n",
    "\n",
    "\n",
    "def test_tuple_sequal():\n",
    "    v0 = relay.Var(\"v0\")\n",
    "    v1 = relay.Var(\"v1\")\n",
    "    v2 = relay.Var(\"v2\")\n",
    "\n",
    "    # unit value is a valid tuple\n",
    "    assert consistent_equal(relay.Tuple([]), relay.Tuple([]))\n",
    "\n",
    "    tup = relay.Tuple([v0, relay.const(2), relay.const(3), relay.Tuple([relay.const(4)])])\n",
    "    same = relay.Tuple([v0, relay.const(2), relay.const(3), relay.Tuple([relay.const(4)])])\n",
    "\n",
    "    assert consistent_equal(tup, same)\n",
    "\n",
    "    # use the eq_map\n",
    "\n",
    "    let_tup = relay.Let(v1, tup, v1)\n",
    "    let_mapped = relay.Let(\n",
    "        v2, relay.Tuple([v0, relay.const(2), relay.const(3), relay.Tuple([relay.const(4)])]), v2\n",
    "    )\n",
    "\n",
    "    assert consistent_equal(let_tup, let_mapped)\n",
    "\n",
    "    more_fields = relay.Tuple(\n",
    "        [v1, relay.const(2), relay.const(3), relay.Tuple([relay.const(4)]), v2]\n",
    "    )\n",
    "    assert not consistent_equal(tup, more_fields)\n",
    "\n",
    "    fewer_fields = relay.Tuple([v1, relay.const(2), relay.const(3)])\n",
    "    assert not consistent_equal(tup, fewer_fields)\n",
    "\n",
    "    different_end = relay.Tuple([v1, relay.const(2), relay.const(3), relay.Tuple([relay.const(5)])])\n",
    "    assert not consistent_equal(tup, different_end)\n",
    "\n",
    "    different_start = relay.Tuple(\n",
    "        [v2, relay.const(2), relay.const(3), relay.Tuple([relay.const(4)])]\n",
    "    )\n",
    "    assert not consistent_equal(tup, different_start)\n",
    "\n",
    "    longer_at_end = relay.Tuple(\n",
    "        [v1, relay.const(2), relay.const(3), relay.Tuple([relay.const(4), relay.const(5)])]\n",
    "    )\n",
    "    assert not consistent_equal(tup, longer_at_end)\n",
    "\n",
    "\n",
    "def test_tuple_get_item_sequal():\n",
    "    x = relay.Var(\"x\")\n",
    "    y = relay.Var(\"y\")\n",
    "    assert not consistent_equal(relay.TupleGetItem(x, 1), relay.TupleGetItem(y, 1))\n",
    "    assert not consistent_equal(relay.TupleGetItem(x, 1), relay.TupleGetItem(x, 2))\n",
    "    assert consistent_equal(relay.TupleGetItem(x, 1), relay.TupleGetItem(x, 1))\n",
    "\n",
    "\n",
    "def test_function_attr():\n",
    "    x0 = relay.var(\"x0\", shape=(10, 10))\n",
    "    w00 = relay.var(\"w00\", shape=(10, 10))\n",
    "    w01 = relay.var(\"w01\", shape=(10, 10))\n",
    "    w02 = relay.var(\"w02\", shape=(10, 10))\n",
    "    z00 = relay.add(x0, w00)\n",
    "    p00 = relay.subtract(z00, w01)\n",
    "    q00 = relay.multiply(p00, w02)\n",
    "    func0 = relay.Function([x0, w00, w01, w02], q00)\n",
    "    func0 = func0.with_attr(\"FuncName\", \"a\")\n",
    "\n",
    "    x1 = relay.var(\"x1\", shape=(10, 10))\n",
    "    w10 = relay.var(\"w10\", shape=(10, 10))\n",
    "    w11 = relay.var(\"w11\", shape=(10, 10))\n",
    "    w12 = relay.var(\"w12\", shape=(10, 10))\n",
    "    z10 = relay.add(x1, w10)\n",
    "    p10 = relay.subtract(z10, w11)\n",
    "    q10 = relay.multiply(p10, w12)\n",
    "    func1 = relay.Function([x1, w10, w11, w12], q10)\n",
    "    func1 = func1.with_attr(\"FuncName\", \"b\")\n",
    "    assert not consistent_equal(func0, func1)\n",
    "\n",
    "\n",
    "def test_function_sequal():\n",
    "    tt1 = relay.TensorType((1, 2, 3), \"float32\")\n",
    "    tt2 = relay.TensorType((4, 5, 6), \"int8\")\n",
    "    tt3 = relay.TupleType([tt1, tt2])\n",
    "\n",
    "    v1 = relay.Var(\"v1\", tt1)\n",
    "    v2 = relay.Var(\"v2\", tt2)\n",
    "    v3 = relay.Var(\"v3\", tt3)\n",
    "    v4 = relay.Var(\"v4\", tt2)\n",
    "    vret = relay.Constant(tvm.nd.array(np.ones(1)))\n",
    "\n",
    "    tp1 = relay.TypeVar(\"tp1\", relay.TypeKind.Type)\n",
    "    tp2 = relay.TypeVar(\"tp2\", relay.TypeKind.Type)\n",
    "    tp3 = relay.TypeVar(\"tp3\", relay.TypeKind.ShapeVar)\n",
    "    tp4 = relay.TypeVar(\"tp4\", relay.TypeKind.ShapeVar)\n",
    "\n",
    "    basic_args = [relay.Var(\"v3\", tt1), relay.Var(\"v4\", tt2)]\n",
    "    basic_tps = [tp1, tp2]\n",
    "\n",
    "    func = relay.Function([v1, v2], v1, tt2, basic_tps)\n",
    "    mapped = relay.Function(basic_args, basic_args[0], tt2, basic_tps)\n",
    "    assert consistent_equal(func, mapped)\n",
    "\n",
    "    fewer_params = relay.Function([relay.Var(\"v4\", tt2)], v4, tt2, basic_tps)\n",
    "    assert not consistent_equal(func, fewer_params)\n",
    "\n",
    "    more_params = relay.Function(\n",
    "        [relay.Var(\"v3\", tt1), relay.Var(\"v4\", tt2), relay.Var(\"v2\", tt2)], v4, tt2, basic_tps\n",
    "    )\n",
    "    assert not consistent_equal(func, more_params)\n",
    "\n",
    "    params_unordered = relay.Function([v2, v1], v1, tt2, basic_tps)\n",
    "    assert not consistent_equal(func, params_unordered)\n",
    "\n",
    "    params_mismatch = relay.Function([v1, v3], v1, tt2, basic_tps)\n",
    "    assert not consistent_equal(func, params_mismatch)\n",
    "\n",
    "    # also would not typecheck\n",
    "    ret_type_mismatch = relay.Function(basic_args, v4, tt1, basic_tps)\n",
    "    assert not consistent_equal(func, ret_type_mismatch)\n",
    "\n",
    "    # also mis-typed\n",
    "    different_body = relay.Function(basic_args, v3, tt2, basic_tps)\n",
    "    assert not consistent_equal(func, different_body)\n",
    "\n",
    "    fewer_type_params = relay.Function(basic_args, v4, tt2, [tp1])\n",
    "    assert not consistent_equal(func, fewer_type_params)\n",
    "\n",
    "    more_type_params = relay.Function(basic_args, v4, tt2, [tp1, tp2, tp3])\n",
    "    assert not consistent_equal(func, more_type_params)\n",
    "\n",
    "    type_params_unordered = relay.Function(basic_args, v4, tt2, [tp2, tp1])\n",
    "    assert not consistent_equal(func, type_params_unordered)\n",
    "\n",
    "    different_type_params = relay.Function(basic_args, v4, tt2, [tp3, tp4])\n",
    "    assert not consistent_equal(func, different_type_params)\n",
    "\n",
    "    # a well-typed example that also differs in body, ret type, and type params\n",
    "    tupled_example = relay.Function(basic_args, relay.Tuple([v3, v4]), tt3)\n",
    "    assert not consistent_equal(func, tupled_example)\n",
    "\n",
    "    # nullable\n",
    "    no_ret_type = relay.Function(basic_args, v4, None, [tp1, tp2])\n",
    "    # both null\n",
    "    assert consistent_equal(no_ret_type, no_ret_type)\n",
    "    # one null\n",
    "    assert not consistent_equal(func, no_ret_type)\n",
    "    assert not consistent_equal(no_ret_type, func)\n",
    "\n",
    "\n",
    "def test_call_sequal():\n",
    "    v1 = relay.Var(\"v1\")\n",
    "    v2 = relay.Var(\"v2\")\n",
    "\n",
    "    attr1 = tvm.ir.make_node(\"attrs.TestAttrs\", name=\"attr\", padding=(3, 4))\n",
    "    attr1_same = tvm.ir.make_node(\"attrs.TestAttrs\", name=\"attr\", padding=(3, 4))\n",
    "    attr2 = tvm.ir.make_node(\"attrs.TestAttrs\", name=\"attr\", padding=(3, 4, 4))\n",
    "\n",
    "    tt1 = relay.TensorType((1, 2, 3), \"float32\")\n",
    "    tt2 = relay.TensorType((), \"int8\")\n",
    "\n",
    "    basic_args = [relay.const(1), relay.const(2), v2, relay.Tuple([])]\n",
    "\n",
    "    # manually writing out args to ensure that args does not rely on\n",
    "    # pointer equality\n",
    "    call = relay.Call(v1, [relay.const(1), relay.const(2), v2, relay.Tuple([])], attr1, [tt1])\n",
    "    same = relay.Call(v1, basic_args, attr1, [tt1])\n",
    "    assert consistent_equal(call, same)\n",
    "\n",
    "    different_fn = relay.Call(v2, basic_args, attr1, [tt1])\n",
    "    assert not consistent_equal(call, different_fn)\n",
    "\n",
    "    fewer_args = relay.Call(v1, [relay.const(1), relay.const(2), v2], attr1, [tt1])\n",
    "    assert not consistent_equal(call, fewer_args)\n",
    "\n",
    "    reordered_args = relay.Call(\n",
    "        v1, [relay.const(2), relay.const(1), relay.Tuple([]), v2], attr1, [tt1]\n",
    "    )\n",
    "    assert not consistent_equal(call, reordered_args)\n",
    "\n",
    "    different_args = relay.Call(v1, [relay.const(1), relay.const(2), relay.const(3)], attr1, [tt1])\n",
    "    assert not consistent_equal(call, different_args)\n",
    "\n",
    "    more_args = relay.Call(\n",
    "        v1,\n",
    "        [relay.const(1), relay.const(2), v2, relay.Tuple([]), relay.const(3), relay.const(4)],\n",
    "        attr1,\n",
    "        [tt1],\n",
    "    )\n",
    "    assert not consistent_equal(call, more_args)\n",
    "\n",
    "    different_attrs = relay.Call(v1, basic_args, attr2, [tt1])\n",
    "    assert not consistent_equal(call, different_attrs)\n",
    "\n",
    "    same_attrs = relay.Call(v1, basic_args, attr1_same, [tt1])\n",
    "    assert consistent_equal(call, same_attrs)\n",
    "\n",
    "    no_type_args = relay.Call(v1, basic_args, attr1)\n",
    "    assert not consistent_equal(call, no_type_args)\n",
    "\n",
    "    more_type_args = relay.Call(v1, basic_args, attr1, [tt1, tt2])\n",
    "    assert not consistent_equal(call, more_type_args)\n",
    "\n",
    "    different_type_arg = relay.Call(v1, basic_args, attr1, [tt2])\n",
    "    assert not consistent_equal(call, different_type_arg)\n",
    "\n",
    "\n",
    "def test_let_sequal():\n",
    "    tt1 = relay.TensorType((), \"float32\")\n",
    "    tt2 = relay.TensorType((), \"int8\")\n",
    "    v1 = relay.Var(\"v1\")\n",
    "    v1_wtype = relay.Var(\"v1\", tt1)\n",
    "    v2 = relay.Var(\"v2\")\n",
    "    v3 = relay.Var(\"v3\")\n",
    "\n",
    "    let = relay.Let(v1, relay.const(2), v1)\n",
    "    mapped = relay.Let(v2, relay.const(2), v2)\n",
    "    assert consistent_equal(let, mapped)\n",
    "\n",
    "    mismatched_var = relay.Let(v2, relay.const(2), v3)\n",
    "    assert not consistent_equal(let, mismatched_var)\n",
    "\n",
    "    different_value = relay.Let(v2, relay.const(3), v2)\n",
    "    assert not consistent_equal(let, different_value)\n",
    "\n",
    "    different_body = relay.Let(v2, relay.const(3), relay.const(12))\n",
    "    assert not consistent_equal(let, different_body)\n",
    "\n",
    "    # specified types must match\n",
    "\n",
    "    let_with_type = relay.Let(v1_wtype, relay.const(2), v1_wtype)\n",
    "    same_type = relay.Let(v1_wtype, relay.const(2), v1_wtype)\n",
    "    assert consistent_equal(let_with_type, same_type)\n",
    "    assert not consistent_equal(let, let_with_type)\n",
    "    v2 = relay.Var(\"v1\", tt2)\n",
    "    different_type = relay.Let(v2, relay.const(2), v2)\n",
    "    assert not consistent_equal(let_with_type, different_type)\n",
    "\n",
    "\n",
    "def test_if_sequal():\n",
    "    v1 = relay.Var(\"v1\")\n",
    "    v2 = relay.Var(\"v2\")\n",
    "\n",
    "    if_sample = relay.If(v1, relay.const(1), relay.Tuple([relay.const(2), relay.const(3)]))\n",
    "    same = relay.If(v1, relay.const(1), relay.Tuple([relay.const(2), relay.const(3)]))\n",
    "    assert consistent_equal(if_sample, same)\n",
    "\n",
    "    different_cond = relay.If(v2, relay.const(1), relay.Tuple([relay.const(2), relay.const(3)]))\n",
    "    assert not consistent_equal(if_sample, different_cond)\n",
    "\n",
    "    different_true = relay.If(v1, relay.const(2), relay.Tuple([relay.const(2), relay.const(3)]))\n",
    "    assert not consistent_equal(if_sample, different_true)\n",
    "\n",
    "    different_false = relay.If(v1, relay.const(1), relay.Tuple([]))\n",
    "    assert not consistent_equal(if_sample, different_false)\n",
    "\n",
    "\n",
    "def test_constructor_sequal():\n",
    "    # smoke test: it should be pointer equality\n",
    "    mod = tvm.IRModule()\n",
    "    p = relay.prelude.Prelude(mod)\n",
    "    _, cons, nil = p.mod.get_type(\"List\")\n",
    "\n",
    "    assert consistent_equal(nil, nil)\n",
    "    assert consistent_equal(cons, cons)\n",
    "    assert not consistent_equal(nil, cons)\n",
    "\n",
    "\n",
    "def test_match_sequal():\n",
    "    mod = tvm.IRModule()\n",
    "    p = relay.prelude.Prelude(mod)\n",
    "    _, cons, nil = p.mod.get_type(\"List\")\n",
    "    _, none, some = p.mod.get_type(\"Option\")\n",
    "\n",
    "    x = relay.Var(\"x\")\n",
    "    y = relay.Var(\"y\")\n",
    "    nil_case = relay.Clause(relay.PatternConstructor(nil), nil())\n",
    "    cons_case = relay.Clause(\n",
    "        relay.PatternConstructor(cons, [relay.PatternVar(x), relay.PatternVar(y)]), cons(x, y)\n",
    "    )\n",
    "\n",
    "    z = relay.Var(\"z\")\n",
    "    a = relay.Var(\"a\")\n",
    "    equivalent_cons = relay.Clause(\n",
    "        relay.PatternConstructor(cons, [relay.PatternVar(z), relay.PatternVar(a)]), cons(z, a)\n",
    "    )\n",
    "\n",
    "    data = cons(relay.const(1), cons(relay.const(2), nil()))\n",
    "\n",
    "    match = relay.Match(data, [nil_case, cons_case])\n",
    "    equivalent = relay.Match(data, [nil_case, equivalent_cons])\n",
    "    empty = relay.Match(data, [])\n",
    "    no_cons = relay.Match(data, [nil_case])\n",
    "    no_nil = relay.Match(data, [cons_case])\n",
    "    different_data = relay.Match(nil(), [nil_case, cons_case])\n",
    "    different_order = relay.Match(data, [cons_case, nil_case])\n",
    "    different_nil = relay.Match(\n",
    "        data, [relay.Clause(relay.PatternConstructor(nil), cons(nil(), nil())), cons_case]\n",
    "    )\n",
    "    different_cons = relay.Match(\n",
    "        data,\n",
    "        [\n",
    "            nil_case,\n",
    "            relay.Clause(\n",
    "                relay.PatternConstructor(cons, [relay.PatternWildcard(), relay.PatternWildcard()]),\n",
    "                nil(),\n",
    "            ),\n",
    "        ],\n",
    "    )\n",
    "    another_case = relay.Match(\n",
    "        data, [nil_case, cons_case, relay.Clause(relay.PatternWildcard(), nil())]\n",
    "    )\n",
    "    wrong_constructors = relay.Match(\n",
    "        data,\n",
    "        [\n",
    "            relay.Clause(relay.PatternConstructor(none), nil()),\n",
    "            relay.Clause(relay.PatternConstructor(some, [relay.PatternVar(x)]), cons(x, nil())),\n",
    "        ],\n",
    "    )\n",
    "\n",
    "    tvm.ir.assert_structural_equal(match, match)\n",
    "    assert consistent_equal(match, match)\n",
    "    assert consistent_equal(match, equivalent)\n",
    "    assert not consistent_equal(match, no_cons)\n",
    "    assert not consistent_equal(match, no_nil)\n",
    "    assert not consistent_equal(match, empty)\n",
    "    assert not consistent_equal(match, different_data)\n",
    "    assert not consistent_equal(match, different_order)\n",
    "    assert not consistent_equal(match, different_nil)\n",
    "    assert not consistent_equal(match, different_cons)\n",
    "    assert not consistent_equal(match, another_case)\n",
    "    assert not consistent_equal(match, wrong_constructors)\n",
    "\n",
    "\n",
    "def test_op_sequal():\n",
    "    # only checks names\n",
    "    op1 = relay.op.get(\"add\")\n",
    "    op2 = relay.op.get(\"add\")\n",
    "    assert consistent_equal(op1, op2)\n",
    "\n",
    "    op3 = relay.op.get(\"take\")\n",
    "    assert not consistent_equal(op1, op3)\n",
    "\n",
    "\n",
    "def test_graph_equal():\n",
    "    x = relay.var(\"x\")\n",
    "\n",
    "    y0 = relay.add(x, x)\n",
    "    z0 = relay.add(y0, y0)\n",
    "\n",
    "    y1 = relay.add(x, x)\n",
    "    z1 = relay.add(y1, y1)\n",
    "\n",
    "    z3 = relay.add(relay.add(x, x), relay.add(x, x))\n",
    "\n",
    "    assert consistent_equal(z0, z1)\n",
    "    assert consistent_equal(z0, z1)\n",
    "\n",
    "    # z3's dataflow format is different from z0\n",
    "    # z0 is computed from a common y0 node\n",
    "    # Relay view them as different programs\n",
    "    # Check the difference in the text format.\n",
    "    assert not consistent_equal(z0, z3)\n",
    "\n",
    "\n",
    "def test_hash_unequal():\n",
    "    x1 = relay.var(\"x1\", shape=(10, 10), dtype=\"float32\")\n",
    "    y1 = relay.var(\"y1\", shape=(10, 10), dtype=\"float32\")\n",
    "    func1 = relay.Function([x1, y1], relay.add(x1, y1))\n",
    "\n",
    "    # func2 is exactly same structure with same variables shapes and dtypes\n",
    "    x2 = relay.var(\"x2\", shape=(10, 10), dtype=\"float32\")\n",
    "    y2 = relay.var(\"y2\", shape=(10, 10), dtype=\"float32\")\n",
    "    func2 = relay.Function([x2, y2], relay.add(x2, y2))\n",
    "\n",
    "    assert consistent_equal(func1, func2)\n",
    "\n",
    "    # func3 is same as func1 but with different var shapes\n",
    "    x3 = relay.var(\"x3\", shape=(20, 10), dtype=\"float32\")\n",
    "    y3 = relay.var(\"y3\", shape=(20, 10), dtype=\"float32\")\n",
    "    func3 = relay.Function([x3, y3], relay.add(x3, y3))\n",
    "\n",
    "    assert not consistent_equal(func1, func3)\n",
    "\n",
    "\n",
    "def test_tuple_match():\n",
    "    a = relay.Var(\"a\")\n",
    "    b = relay.Var(\"b\")\n",
    "    clause = relay.Clause(relay.PatternTuple([relay.PatternVar(a), relay.PatternVar(b)]), a + b)\n",
    "    x = relay.Match(relay.Tuple([relay.const(1), relay.const(1)]), [clause])\n",
    "\n",
    "    a = relay.Var(\"a\")\n",
    "    b = relay.Var(\"b\")\n",
    "    clause = relay.Clause(relay.PatternTuple([relay.PatternVar(a), relay.PatternVar(b)]), a + b)\n",
    "    y = relay.Match(relay.Tuple([relay.const(1), relay.const(1)]), [clause])\n",
    "    assert consistent_equal(x, y)\n",
    "\n",
    "\n",
    "def test_fn_attribute():\n",
    "    # create function that performs add\n",
    "    a = relay.var(\"a\", shape=(10, 10))\n",
    "    b = relay.var(\"b\", shape=(10, 10))\n",
    "    add = relay.add(a, b)\n",
    "    add_fn = relay.Function([a, b], add)\n",
    "    add_fn = run_opt_pass(add_fn, relay.transform.InferType())\n",
    "\n",
    "    # create function that performs add with test attribute\n",
    "    c = relay.var(\"c\", shape=(10, 10))\n",
    "    d = relay.var(\"d\", shape=(10, 10))\n",
    "    add_1 = relay.add(c, d)\n",
    "    add_1_fn = relay.Function([c, d], add_1)\n",
    "    add_1_fn = add_1_fn.with_attr(\"TestAttribute\", \"test\")\n",
    "    add_1_fn = run_opt_pass(add_1_fn, relay.transform.InferType())\n",
    "\n",
    "    assert not consistent_equal(add_1_fn, add_fn)\n",
    "    assert not consistent_equal(add_fn, add_1_fn)\n",
    "\n",
    "\n",
    "def test_fn_vid_map():\n",
    "    def get_fn(with_vid):\n",
    "        x = relay.var(\"x\", shape=(10,), dtype=\"float32\")\n",
    "        f = relay.Function([x], x).with_attr(\"dict\", {x.vid: 1} if with_vid else {x: 1})\n",
    "        return f\n",
    "\n",
    "    assert consistent_equal(get_fn(True), get_fn(True))\n",
    "    assert consistent_equal(get_fn(False), get_fn(False))\n",
    "\n",
    "\n",
    "def test_lets():\n",
    "    shape = (5, 5)\n",
    "\n",
    "    def func1():\n",
    "        sb = relay.ScopeBuilder()\n",
    "        p0 = relay.var(\"p0\", shape=shape)\n",
    "        p1 = relay.var(\"p1\", shape=shape)\n",
    "        a0 = sb.let(\"a0\", relay.add(p0, relay.const(1)))\n",
    "        a1 = sb.let(\"a1\", relay.add(p1, relay.const(1)))\n",
    "        a2 = sb.let(\"a2\", relay.add(a0, a1))\n",
    "        sb.ret(a2)\n",
    "        return relay.Function([p0, p1], sb.get())\n",
    "\n",
    "    def func2():\n",
    "        # Alpha conversion is structurally equal\n",
    "        sb = relay.ScopeBuilder()\n",
    "        p0 = relay.var(\"p0\", shape=shape)\n",
    "        p1 = relay.var(\"p1\", shape=shape)\n",
    "        a1 = sb.let(\"a1\", relay.add(p0, relay.const(1)))\n",
    "        a0 = sb.let(\"a0\", relay.add(p1, relay.const(1)))\n",
    "        a2 = sb.let(\"a2\", relay.add(a1, a0))\n",
    "        sb.ret(a2)\n",
    "        return relay.Function([p0, p1], sb.get())\n",
    "\n",
    "    def func3():\n",
    "        # But changing the order of bindings is not structurally equal\n",
    "        # (even though algebraically equal)\n",
    "        sb = relay.ScopeBuilder()\n",
    "        p0 = relay.var(\"p0\", shape=shape)\n",
    "        p1 = relay.var(\"p1\", shape=shape)\n",
    "        a1 = sb.let(\"a1\", relay.add(p1, relay.const(1)))\n",
    "        a0 = sb.let(\"a0\", relay.add(p0, relay.const(1)))\n",
    "        a2 = sb.let(\"a2\", relay.add(a1, a0))\n",
    "        sb.ret(a2)\n",
    "        return relay.Function([p0, p1], sb.get())\n",
    "\n",
    "    assert tvm.ir.structural_equal(func1(), func2())\n",
    "    assert not tvm.ir.structural_equal(func1(), func3())\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py312x",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
